import xarray as xr
import numpy as np
import os
# 运行目录需要放在根目录
# 变量映射
var_map = {
    "t2m": "2m_temperature",
    "u10": "10m_u_component_of_wind",
    "v10": "10m_v_component_of_wind"
}

# 输入文件路径
input_nc = "era5_data/surface_2023-01-19.nc"  # 你的一天的 ERA5 文件
output_dir = "obs"
os.makedirs(output_dir, exist_ok=True)

# 从哪个编号开始输出
start_index = 72# 如果你前面已经生成了一些，可以设置成 24, 48, ...

# 打开 ERA5 文件
ds = xr.open_dataset(input_nc)

for hour in range(24):
    out_vars = {}

    for old_name, new_name in var_map.items():
        var = ds[old_name].isel(valid_time=hour)  # 取第 hour 小时

        # 如果没有时间维，加一维 (1, H, W)
        var = var.expand_dims("time")

        # 下采样到 60x120
        var_ds = var.coarsen(latitude=12, longitude=12, boundary='trim').mean()
        var_ds = var_ds.rename({
            "latitude": "height",
            "longitude": "width"
        })

        out_vars[new_name] = var_ds

    # 合成新的数据集
    new_ds = xr.Dataset(out_vars)
    new_ds = new_ds.assign_coords({
        "time": [0],
        "height": np.arange(60),
        "width": np.arange(120)
    })

    # 保存文件
    output_path = os.path.join(output_dir, f"obs_{start_index + hour:03d}.nc")
    new_ds.to_netcdf(output_path)
    print(f"✅ 保存 {output_path}")
